#include "stdafx.h"

#include <windows.h>
#include "RegKey.h"

using namespace WCMM;

CRegKey::CRegKey() throw() :
	m_hKey( NULL ), m_samWOW64(0)
{
}
CRegKey::CRegKey(_Inout_ CRegKey& key) throw() :
	m_hKey( NULL )
{
	REGSAM samWOW64 = key.m_samWOW64;
	Attach( key.Detach() );
	m_samWOW64 = samWOW64;
}

CRegKey::CRegKey(_In_ HKEY hKey) throw() :
	m_hKey(hKey), m_samWOW64(0)
{
}

CRegKey::~CRegKey() throw()
{
	Close();
}

CRegKey& CRegKey::operator=(_Inout_ CRegKey& key) throw()
{
	if(m_hKey!=key.m_hKey){
		Close();
		REGSAM samWOW64 = key.m_samWOW64;
		Attach( key.Detach() );
		m_samWOW64 = samWOW64;
	}
	return( *this );
}

CRegKey::operator HKEY() const throw()
{
	return m_hKey;
}

HKEY CRegKey::Detach() throw()
{
	HKEY hKey = m_hKey;
	m_hKey = NULL;
	m_samWOW64 = 0;
	return hKey;
}

void CRegKey::Attach(_In_ HKEY hKey) throw()
{
	ASSERTE(m_hKey == NULL);
	m_hKey = hKey;
	m_samWOW64 = 0;
}

LONG CRegKey::DeleteSubKey(_In_z_ LPCTSTR lpszSubKey) throw()
{
	ASSERTE(m_hKey != NULL);

#if WINVER >= 0x0501
#ifdef _UNICODE
	static decltype(RegDeleteKeyExW) *pfnRegDeleteKeyEx = NULL;
#else
	static decltype(RegDeleteKeyExA) *pfnRegDeleteKeyEx = NULL;
#endif
	static bool bInitialized = false;

	if (!bInitialized){
		HMODULE hAdvapi32 = GetModuleHandle(_T("Advapi32.dll"));
		if (hAdvapi32 != NULL){
#ifdef _UNICODE
			pfnRegDeleteKeyEx = (decltype(RegDeleteKeyExW) *)GetProcAddress(hAdvapi32, "RegDeleteKeyExW");
#else
			pfnRegDeleteKeyEx = (decltype(RegDeleteKeyExA) *)GetProcAddress(hAdvapi32, "RegDeleteKeyExA");
#endif
		}
		bInitialized = true;
	}

	if (pfnRegDeleteKeyEx != NULL){
		return pfnRegDeleteKeyEx(m_hKey, lpszSubKey, m_samWOW64, 0);
	}

#endif

	return RegDeleteKey(m_hKey, lpszSubKey);
}

LONG CRegKey::DeleteValue(_In_z_ LPCTSTR lpszValue) throw()
{
	ASSERTE(m_hKey != NULL);
	return RegDeleteValue(m_hKey, (LPTSTR)lpszValue);
}

LONG CRegKey::Close() throw()
{
	LONG lRes = ERROR_SUCCESS;
	if (m_hKey != NULL){
		lRes = RegCloseKey(m_hKey);
		m_hKey = NULL;
	}
	m_samWOW64 = 0;
	return lRes;
}

LONG CRegKey::Flush() throw()
{
	ASSERTE(m_hKey != NULL);

	return ::RegFlushKey(m_hKey);
}

LONG CRegKey::EnumKey(
	_In_ DWORD iIndex,
	_Out_writes_to_(*pnNameLength, *pnNameLength) _Post_z_ LPTSTR pszName,
	_Inout_ LPDWORD pnNameLength,
	_Out_opt_ FILETIME* pftLastWriteTime) throw()
{
	FILETIME ftLastWriteTime;

	ASSERTE(m_hKey != NULL);
	if (pftLastWriteTime == NULL){
		pftLastWriteTime = &ftLastWriteTime;
	}

	return ::RegEnumKeyEx(m_hKey, iIndex, pszName, pnNameLength, NULL, NULL, NULL, pftLastWriteTime);
}

LONG CRegKey::NotifyChangeKeyValue(
	_In_ BOOL bWatchSubtree,
	_In_ DWORD dwNotifyFilter,
	_In_ HANDLE hEvent,
	_In_ BOOL bAsync) throw()
{
	ASSERTE(m_hKey != NULL);
	ASSERTE((hEvent != NULL) || !bAsync);

	return ::RegNotifyChangeKeyValue(m_hKey, bWatchSubtree, dwNotifyFilter, hEvent, bAsync);
}

LONG CRegKey::Create(
	_In_ HKEY hKeyParent,
	_In_z_ LPCTSTR lpszKeyName,
	_In_opt_z_ LPTSTR lpszClass,
	_In_ DWORD dwOptions,
	_In_ REGSAM samDesired,
	_In_opt_ LPSECURITY_ATTRIBUTES lpSecAttr,
	_Out_opt_ LPDWORD lpdwDisposition) throw()
{
	ASSERTE(hKeyParent != NULL);

	DWORD dw;
	HKEY hKey = NULL;
	LONG lRes = RegCreateKeyEx(hKeyParent, lpszKeyName, 0, lpszClass, dwOptions, samDesired, lpSecAttr, &hKey, &dw);
	if (lpdwDisposition != NULL)
		*lpdwDisposition = dw;
	if (lRes == ERROR_SUCCESS){
		lRes = Close();
		m_hKey = hKey;
#if WINVER >= 0x0501
		m_samWOW64 = samDesired & (KEY_WOW64_32KEY | KEY_WOW64_64KEY);
#endif
	}
	return lRes;
}

LONG CRegKey::Open(
	_In_ HKEY hKeyParent,
	_In_opt_z_ LPCTSTR lpszKeyName,
	_In_ REGSAM samDesired) throw()
{
	ASSERTE(hKeyParent != NULL);
	HKEY hKey = NULL;
	LONG lRes = RegOpenKeyEx(hKeyParent, lpszKeyName, 0, samDesired, &hKey);
	if (lRes == ERROR_SUCCESS){
		lRes = Close();
		m_hKey = hKey;
#if WINVER >= 0x0501
		m_samWOW64 = samDesired & (KEY_WOW64_32KEY | KEY_WOW64_64KEY);
#endif
	}
	return lRes;
}

LONG CRegKey::QueryValue(
	_Out_ DWORD& dwValue,
	_In_opt_z_ LPCTSTR lpszValueName)
{
	DWORD dwType = 0;
	DWORD dwCount = sizeof(DWORD);
	LONG lRes = RegQueryValueEx(m_hKey, lpszValueName, NULL, &dwType,
		(LPBYTE)&dwValue, &dwCount);
	ASSERTE((lRes!=ERROR_SUCCESS) || (dwType == REG_DWORD));
	ASSERTE((lRes!=ERROR_SUCCESS) || (dwCount == sizeof(DWORD)));
	if (dwType != REG_DWORD)
		return ERROR_INVALID_DATA;
	return lRes;
}

LONG CRegKey::QueryValue(
	_Out_writes_to_opt_(*pdwCount, *pdwCount) LPTSTR pszValue,
	_In_opt_z_ LPCTSTR lpszValueName,
	_Inout_ DWORD* pdwCount)
{
	ASSERTE(pdwCount != NULL);
	DWORD dwType = 0;
	LONG lRes = RegQueryValueEx(m_hKey, lpszValueName, NULL, &dwType, (LPBYTE)pszValue, pdwCount);
	ASSERTE((lRes!=ERROR_SUCCESS) || (dwType == REG_SZ) ||
		(dwType == REG_MULTI_SZ) || (dwType == REG_EXPAND_SZ));
	if (pszValue != NULL){
		if(*pdwCount>0){
			switch(dwType){
			case REG_SZ:
			case REG_EXPAND_SZ:
				if ((*pdwCount) % sizeof(TCHAR) != 0 || pszValue[(*pdwCount) / sizeof(TCHAR) - 1] != 0){
					pszValue[0]=_T('\0');
					return ERROR_INVALID_DATA;
				}
				break;
			case REG_MULTI_SZ:
				if ((*pdwCount) % sizeof(TCHAR) != 0 || (*pdwCount) / sizeof(TCHAR) < 1 || 
					pszValue[(*pdwCount) / sizeof(TCHAR) -1] != 0 || 
					(((*pdwCount) / sizeof(TCHAR))>1 && pszValue[(*pdwCount) / sizeof(TCHAR) - 2] != 0) ){
						pszValue[0]=_T('\0');
						return ERROR_INVALID_DATA;
				}
				break;
			default:

				//
				// Ensure termination
				//

				pszValue[0]=_T('\0');
				return ERROR_INVALID_DATA;
			}
		}else{

			//
			// this is a blank one with no data yet
			// Ensure termination
			//

			pszValue[0]=_T('\0');
		}
	}
	return lRes;
}

LONG CRegKey::QueryValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_Out_opt_ DWORD* pdwType,
	_Out_opt_ void* pData,
	_Inout_ ULONG* pnBytes) throw()
{
	ASSERTE(m_hKey != NULL);

	return( ::RegQueryValueEx(m_hKey, pszValueName, NULL, pdwType, static_cast< LPBYTE >( pData ), pnBytes) );
}

LONG CRegKey::QueryDWORDValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_Out_ DWORD& dwValue) throw()
{
	LONG lRes;
	ULONG nBytes;
	DWORD dwType;

	ASSERTE(m_hKey != NULL);

	nBytes = sizeof(DWORD);
	lRes = ::RegQueryValueEx(m_hKey, pszValueName, NULL, &dwType, reinterpret_cast<LPBYTE>(&dwValue),
		&nBytes);
	if (lRes != ERROR_SUCCESS)
		return lRes;
	if (dwType != REG_DWORD)
		return ERROR_INVALID_DATA;

	return ERROR_SUCCESS;
}
LONG CRegKey::QueryQWORDValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_Out_ ULONGLONG& qwValue) throw()
{
	LONG lRes;
	ULONG nBytes;
	DWORD dwType;

	ASSERTE(m_hKey != NULL);

	nBytes = sizeof(ULONGLONG);
	lRes = ::RegQueryValueEx(m_hKey, pszValueName, NULL, &dwType, reinterpret_cast<LPBYTE>(&qwValue),
		&nBytes);
	if (lRes != ERROR_SUCCESS)
		return lRes;
	if (dwType != REG_QWORD)
		return ERROR_INVALID_DATA;

	return ERROR_SUCCESS;
}

LONG CRegKey::QueryBinaryValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_Out_opt_ void* pValue,
	_Inout_opt_ ULONG* pnBytes) throw()
{
	LONG lRes;
	DWORD dwType;

	ASSERTE(pnBytes != NULL);
	ASSERTE(m_hKey != NULL);

	lRes = ::RegQueryValueEx(m_hKey, pszValueName, NULL, &dwType, reinterpret_cast<LPBYTE>(pValue),
		pnBytes);
	if (lRes != ERROR_SUCCESS)
		return lRes;
	if (dwType != REG_BINARY)
		return ERROR_INVALID_DATA;

	return ERROR_SUCCESS;
}

LONG CRegKey::QueryStringValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_Out_writes_to_opt_(*pnChars, *pnChars) LPTSTR pszValue,
	_Inout_ ULONG* pnChars) throw()
{
	LONG lRes;
	DWORD dwType;
	ULONG nBytes;

	ASSERTE(m_hKey != NULL);
	ASSERTE(pnChars != NULL);

	nBytes = (*pnChars)*sizeof(TCHAR);
	*pnChars = 0;
	lRes = ::RegQueryValueEx(m_hKey, pszValueName, NULL, &dwType, reinterpret_cast<LPBYTE>(pszValue),
		&nBytes);

	if (lRes != ERROR_SUCCESS){
		return lRes;
	}

	if(dwType != REG_SZ && dwType != REG_EXPAND_SZ){
		return ERROR_INVALID_DATA;
	}

	if (pszValue != NULL){
		if(nBytes!=0){
			if ((nBytes % sizeof(TCHAR) != 0) || (pszValue[nBytes / sizeof(TCHAR) -1] != 0)){
				return ERROR_INVALID_DATA;
			}
		}else{
			pszValue[0]=_T('\0');
		}
	}

	*pnChars = nBytes/sizeof(TCHAR);

	return ERROR_SUCCESS;
}

LONG CRegKey::QueryMultiStringValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_Out_writes_to_opt_(*pnChars, *pnChars) LPTSTR pszValue,
	_Inout_ ULONG* pnChars) throw()
{
	LONG lRes;
	DWORD dwType;
	ULONG nBytes;

	ASSERTE(m_hKey != NULL);
	ASSERTE(pnChars != NULL);

	if (pszValue != NULL && *pnChars < 2)
		return ERROR_INSUFFICIENT_BUFFER;

	nBytes = (*pnChars)*sizeof(TCHAR);
	*pnChars = 0;

	lRes = ::RegQueryValueEx(m_hKey, pszValueName, NULL, &dwType, reinterpret_cast<LPBYTE>(pszValue),
		&nBytes);
	if (lRes != ERROR_SUCCESS)
		return lRes;

	if (dwType != REG_MULTI_SZ)
		return ERROR_INVALID_DATA;

	if (pszValue != NULL && (nBytes % sizeof(TCHAR) != 0 || nBytes / sizeof(TCHAR) < 1 || pszValue[nBytes / sizeof(TCHAR) -1] != 0 || ((nBytes/sizeof(TCHAR))>1 && pszValue[nBytes / sizeof(TCHAR) - 2] != 0)))
		return ERROR_INVALID_DATA;

	*pnChars = nBytes/sizeof(TCHAR);

	return ERROR_SUCCESS;
}

// LONG CWMRegKey::QueryGUIDValue(
// 	_In_opt_z_ LPCTSTR pszValueName,
// 	_Out_ GUID& guidValue) throw()
// {
// 	TCHAR szGUID[64];
// 	LONG lRes;
// 	ULONG nCount;
// 	HRESULT hr;
// 
// 	ASSERTE(m_hKey != NULL);
// 
// 	guidValue = GUID_NULL;
// 
// 	nCount = 64;
// 	lRes = QueryStringValue(pszValueName, szGUID, &nCount);
// 
// 	if (lRes != ERROR_SUCCESS)
// 		return lRes;
// 
// 	if(szGUID[0] != _T('{'))
// 		return ERROR_INVALID_DATA;
// 
// 	USES_CONVERSION_EX;
// 	LPOLESTR lpstr = T2OLE_EX(szGUID, _ATL_SAFE_ALLOCA_DEF_THRESHOLD);
// #ifndef _UNICODE
// 	if(lpstr == NULL)
// 		return E_OUTOFMEMORY;
// #endif
// 
// 	hr = ::CLSIDFromString(lpstr, &guidValue);
// 	if (FAILED(hr))
// 		return ERROR_INVALID_DATA;
// 
// 	return ERROR_SUCCESS;
// }

LONG WINAPI CRegKey::SetValue(
	_In_ HKEY hKeyParent,
	_In_z_ LPCTSTR lpszKeyName,
	_In_opt_z_ LPCTSTR lpszValue,
	_In_opt_z_ LPCTSTR lpszValueName)
{
	ASSERTE(lpszValue != NULL);
	CRegKey key;
	LONG lRes = key.Create(hKeyParent, lpszKeyName);
	if (lRes == ERROR_SUCCESS)
		lRes = key.SetStringValue(lpszValueName, lpszValue);
	return lRes;
}

LONG CRegKey::SetKeyValue(
	_In_z_ LPCTSTR lpszKeyName,
	_In_opt_z_ LPCTSTR lpszValue,
	_In_opt_z_ LPCTSTR lpszValueName) throw()
{
	ASSERTE(lpszValue != NULL);
	CRegKey key;
	LONG lRes = key.Create(m_hKey, lpszKeyName, REG_NONE, REG_OPTION_NON_VOLATILE, KEY_READ | KEY_WRITE | m_samWOW64);
	if (lRes == ERROR_SUCCESS)
		lRes = key.SetStringValue(lpszValueName, lpszValue);
	return lRes;
}

LONG CRegKey::SetValue(
	_In_ DWORD dwValue,
	_In_opt_z_ LPCTSTR pszValueName)
{
	ASSERTE(m_hKey != NULL);
	return SetDWORDValue(pszValueName, dwValue);
}

LONG CRegKey::SetValue(
	_In_z_ LPCTSTR lpszValue,
	_In_opt_z_ LPCTSTR lpszValueName,
	_In_ bool bMulti,
	_In_ int nValueLen)
{
	//ATLENSURE(lpszValue != NULL);
	ASSERTE(m_hKey != NULL);

	if (bMulti && nValueLen == -1)
		return ERROR_INVALID_PARAMETER;

	if (nValueLen == -1)
		nValueLen = static_cast<int>(lstrlen(lpszValue) + 1);

	DWORD dwType = bMulti ? REG_MULTI_SZ : REG_SZ;

	return ::RegSetValueEx(m_hKey, lpszValueName, 0, dwType,
		reinterpret_cast<const BYTE*>(lpszValue), nValueLen*sizeof(TCHAR));
}

LONG CRegKey::SetValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_In_ DWORD dwType,
	_In_opt_ const void* pValue,
	_In_ ULONG nBytes) throw()
{
	ASSERTE(m_hKey != NULL);
	return ::RegSetValueEx(m_hKey, pszValueName, 0, dwType, static_cast<const BYTE*>(pValue), nBytes);
}

LONG CRegKey::SetBinaryValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_In_opt_ const void* pData,
	_In_ ULONG nBytes) throw()
{
	ASSERTE(m_hKey != NULL);
	return ::RegSetValueEx(m_hKey, pszValueName, 0, REG_BINARY, reinterpret_cast<const BYTE*>(pData), nBytes);
}

LONG CRegKey::SetDWORDValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_In_ DWORD dwValue) throw()
{
	ASSERTE(m_hKey != NULL);
	return ::RegSetValueEx(m_hKey, pszValueName, 0, REG_DWORD, reinterpret_cast<const BYTE*>(&dwValue), sizeof(DWORD));
}

LONG CRegKey::SetQWORDValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_In_ ULONGLONG qwValue) throw()
{
	ASSERTE(m_hKey != NULL);
	return ::RegSetValueEx(m_hKey, pszValueName, 0, REG_QWORD, reinterpret_cast<const BYTE*>(&qwValue), sizeof(ULONGLONG));
}

LONG CRegKey::SetStringValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_In_opt_z_ LPCTSTR pszValue,
	_In_ DWORD dwType) throw()
{
	ASSERTE(m_hKey != NULL);
	WCMENSURE_RETURN_VAL(pszValue != NULL, ERROR_INVALID_DATA);
	ASSERTE((dwType == REG_SZ) || (dwType == REG_EXPAND_SZ));

	return ::RegSetValueEx(m_hKey, pszValueName, 0, dwType, reinterpret_cast<const BYTE*>(pszValue), (static_cast<DWORD>(_tcslen(pszValue))+1)*sizeof(TCHAR));
}

LONG CRegKey::SetMultiStringValue(
	_In_opt_z_ LPCTSTR pszValueName,
	_In_z_ LPCTSTR pszValue) throw()
{
	LPCTSTR pszTemp;
	ULONG nBytes;
	ULONG nLength;

	ASSERTE(m_hKey != NULL);
	WCMENSURE_RETURN_VAL(pszValue != NULL, ERROR_INVALID_DATA);

	//
	// Find the total length (in bytes) of all of the strings, including the
	// terminating '\0' of each string, and the second '\0' that terminates
	// the list.
	//

	nBytes = 0;
	pszTemp = pszValue;
	do{
		nLength = static_cast<ULONG>(lstrlen(pszTemp))+1;
		pszTemp += nLength;
		nBytes += nLength*sizeof(TCHAR);
	} while (nLength != 1);

	return ::RegSetValueEx(m_hKey, pszValueName, 0, REG_MULTI_SZ, reinterpret_cast<const BYTE*>(pszValue),
		nBytes);
}

// LONG CWMRegKey::SetGUIDValue(
// 	_In_opt_z_ LPCTSTR pszValueName,
// 	_In_ REFGUID guidValue) throw()
// {
// 	OLECHAR szGUID[64];
// 
// 	ASSERTE(m_hKey != NULL);
// 
// 	::StringFromGUID2(guidValue, szGUID, 64);
// 
// 	USES_CONVERSION_EX;
// 	LPCTSTR lpstr = OLE2CT_EX(szGUID, _ATL_SAFE_ALLOCA_DEF_THRESHOLD);
// #ifndef _UNICODE
// 	if(lpstr == NULL)
// 		return E_OUTOFMEMORY;
// #endif
// 	return SetStringValue(pszValueName, lpstr);
// }

LONG CRegKey::GetKeySecurity(
	_In_ SECURITY_INFORMATION si,
	_Out_opt_ PSECURITY_DESCRIPTOR psd,
	_Inout_ LPDWORD pnBytes) throw()
{
	ASSERTE(m_hKey != NULL);
	ASSERTE(pnBytes != NULL);

	return ::RegGetKeySecurity(m_hKey, si, psd, pnBytes);
}

LONG CRegKey::SetKeySecurity(
	_In_ SECURITY_INFORMATION si,
	_In_ PSECURITY_DESCRIPTOR psd) throw()
{
	ASSERTE(m_hKey != NULL);
	ASSERTE(psd != NULL);

	return ::RegSetKeySecurity(m_hKey, si, psd);
}

LONG CRegKey::RecurseDeleteKey(_In_z_ LPCTSTR lpszKey) throw()
{
	CRegKey key;
	LONG lRes = key.Open(m_hKey, lpszKey, KEY_READ | KEY_WRITE | m_samWOW64);
	if (lRes != ERROR_SUCCESS){
		if (lRes != ERROR_FILE_NOT_FOUND && lRes != ERROR_PATH_NOT_FOUND){

			//ATLTRACE(atlTraceCOM, 0, _T("CRegKey::RecurseDeleteKey : Failed to Open Key %s(Error = %d)\n"), lpszKey, lRes);

			//
			// TODO TRACE
			//
		}
		return lRes;
	}
	FILETIME time;
	DWORD dwSize = 256;
	TCHAR szBuffer[256];
	while (RegEnumKeyEx(key.m_hKey, 0, szBuffer, &dwSize, NULL, NULL, NULL,
		&time) == ERROR_SUCCESS){

			lRes = key.RecurseDeleteKey(szBuffer);
			if (lRes != ERROR_SUCCESS)
				return lRes;
			dwSize = 256;
	}
	key.Close();
	return DeleteSubKey(lpszKey);
}